import cv2
import numpy as np
import torch
from typing import Dict, List
import json

class ProfessionalLEDDegradation:
    """基于实际LED工程经验的退化模拟器"""
    
    def __init__(self, config_path: str = "degradation_config.json"):
        # 加载基于实际调试经验的配置
        self.config = self.load_engineering_config(config_path)
        
        # 基于你的五高两低经验定义退化类型
        self.degradation_functions = {
            'low_gray_color_shift': self.low_gray_color_shift,
            'non_uniform_brightness': self.non_uniform_brightness,
            'motion_tailing': self.motion_tailing,  # 基于高刷新率经验
            'ic_coupling_artifacts': self.ic_coupling_artifacts,
            'bit_depth_quantization': self.bit_depth_quantization,  # 基于灰度级数经验
            'color_temperature_shift': self.color_temperature_shift  # 基于色准控制经验
        }
    
    def load_engineering_config(self, config_path: str) -> Dict:
        """加载基于实际工程经验的配置"""
        default_config = {
            # 低灰偏色参数（基于低灰校正经验）
            'low_gray': {
                'current_levels': [12, 24, 48, 96],  # 实际电流档位
                'color_shift_intensity': [0.1, 0.3],  # 偏色强度范围
                'gamma_curve_distortion': 0.15,  # Gamma曲线失真
            },
            # 均匀性参数（基于均匀性校正经验）
            'uniformity': {
                'brightness_variation': 0.25,  # 亮度变化率
                'corner_attenuation': 0.4,     # 边角衰减
                'hotspot_intensity': 0.15,     # 热点强度
            },
            # 运动拖尾参数（基于高刷新率经验）
            'motion': {
                'exposure_time': [1/60, 1/120],  # 曝光时间范围
                'blur_kernel_sizes': [7, 15, 23],  # 模糊核尺寸
            },
            # IC耦合参数（基于驱动芯片经验）
            'ic_coupling': {
                'crosstalk_level': 0.08,       # 串扰水平
                'signal_interference': 0.12,   # 信号干扰
            }
        }
        
        try:
            with open(config_path, 'r') as f:
                user_config = json.load(f)
                # 合并配置
                default_config.update(user_config)
        except FileNotFoundError:
            print(f"配置文件 {config_path} 不存在，使用默认工程配置")
        
        return default_config
    
    def low_gray_color_shift(self, image: np.ndarray, current_level: int = 12) -> np.ndarray:
        """模拟低灰偏色 - 基于你的低灰校正经验"""
        # 转换为浮点计算
        img_float = image.astype(np.float32) / 255.0
        
        # 基于电流档位计算偏色强度（小电流偏色更严重）
        intensity = 0.5 - (current_level / 255.0) * 0.4
        
        # 模拟低灰偏色（通常偏红或偏绿）
        color_shift = np.random.choice(['red', 'green'])
        
        if color_shift == 'red':
            # 红色通道增强，蓝色通道减弱
            img_float[:, :, 2] = np.clip(img_float[:, :, 2] + intensity * 0.3, 0, 1)
            img_float[:, :, 0] = np.clip(img_float[:, :, 0] - intensity * 0.2, 0, 1)
        else:
            # 绿色通道增强，红色通道减弱
            img_float[:, :, 1] = np.clip(img_float[:, :, 1] + intensity * 0.25, 0, 1)
            img_float[:, :, 2] = np.clip(img_float[:, :, 2] - intensity * 0.15, 0, 1)
        
        # Gamma曲线失真模拟
        gamma_distorted = np.power(img_float, 1.0 + self.config['low_gray']['gamma_curve_distortion'])
        
        return (np.clip(gamma_distorted, 0, 1) * 255).astype(np.uint8)
    
    def non_uniform_brightness(self, image: np.ndarray) -> np.ndarray:
        """模拟亮度不均匀性 - 基于你的均匀性校正经验"""
        h, w = image.shape[:2]
        
        # 创建亮度不均匀性映射
        # 1. 中心亮边缘暗
        y_coords, x_coords = np.ogrid[:h, :w]
        center_y, center_x = h / 2, w / 2
        
        # 计算到中心的距离
        dist_from_center = np.sqrt((x_coords - center_x)**2 + (y_coords - center_y)**2)
        max_dist = np.sqrt(center_x**2 + center_y**2)
        
        # 中心亮度最高，向边缘衰减
        center_brightness = 1.0
        edge_brightness = 1.0 - self.config['uniformity']['brightness_variation']
        brightness_map = center_brightness - (dist_from_center / max_dist) * (center_brightness - edge_brightness)
        
        # 2. 添加随机热点
        hotspot_mask = np.random.random((h, w)) < 0.1  # 10%的区域有热点
        brightness_map[hotspot_mask] += self.config['uniformity']['hotspot_intensity']
        
        # 3. 边角额外衰减
        corner_mask = (dist_from_center > max_dist * 0.7)
        brightness_map[corner_mask] *= (1 - self.config['uniformity']['corner_attenuation'])
        
        # 应用亮度映射
        result = image.astype(np.float32) * brightness_map[:, :, np.newaxis]
        return np.clip(result, 0, 255).astype(np.uint8)
    
    def motion_tailing(self, image: np.ndarray) -> np.ndarray:
        """模拟运动拖尾 - 基于你的高刷新率经验"""
        # 随机选择运动方向
        direction = np.random.choice(['horizontal', 'vertical'])
        kernel_size = np.random.choice(self.config['motion']['blur_kernel_sizes'])
        
        if direction == 'horizontal':
            # 水平运动模糊
            kernel = np.zeros((kernel_size, kernel_size))
            center = kernel_size // 2
            kernel[center, :] = 1.0 / kernel_size
        else:
            # 垂直运动模糊
            kernel = np.zeros((kernel_size, kernel_size))
            center = kernel_size // 2
            kernel[:, center] = 1.0 / kernel_size
        
        # 应用运动模糊
        blurred = cv2.filter2D(image, -1, kernel)
        
        return blurred
    
    def ic_coupling_artifacts(self, image: np.ndarray) -> np.ndarray:
        """模拟IC耦合伪影 - 基于你的驱动芯片经验"""
        h, w = image.shape[:2]
        result = image.copy().astype(np.float32)
        
        # 模拟串扰（crosstalk）
        # 相邻像素间的颜色渗漏
        crosstalk_strength = self.config['ic_coupling']['crosstalk_level']
        
        # 创建串扰核
        crosstalk_kernel = np.array([
            [0, crosstalk_strength, 0],
            [crosstalk_strength, 1 - 4 * crosstalk_strength, crosstalk_strength],
            [0, crosstalk_strength, 0]
        ])
        
        # 分别对每个通道应用串扰
        for channel in range(3):
            result[:, :, channel] = cv2.filter2D(result[:, :, channel], -1, crosstalk_kernel)
        
        # 模拟信号干扰（随机噪声模式）
        interference = np.random.normal(0, self.config['ic_coupling']['signal_interference'] * 10, (h, w, 3))
        result += interference
        
        return np.clip(result, 0, 255).astype(np.uint8)
    
    def bit_depth_quantization(self, image: np.ndarray, target_bits: int = 6) -> np.ndarray:
        """模拟低比特深度量化 - 基于你的灰度级数经验"""
        # 计算量化步长
        max_val = 2**target_bits - 1
        scale = 255.0 / max_val
        
        # 量化到低比特深度
        quantized = (image / scale).astype(np.uint8)
        # 重新缩放回8比特（引入量化误差）
        restored = (quantized * scale).astype(np.uint8)
        
        return restored
    
    def color_temperature_shift(self, image: np.ndarray) -> np.ndarray:
        """模拟色温偏移 - 基于你的色准控制经验"""
        # 随机选择色温偏移方向（暖色或冷色）
        shift_direction = np.random.choice(['warm', 'cool'])
        
        # 转换为LAB颜色空间进行更精确的颜色调整
        lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
        lab_float = lab.astype(np.float32)
        
        if shift_direction == 'warm':
            # 向暖色偏移：增加红色，减少蓝色
            lab_float[:, :, 2] += np.random.uniform(5, 15)  # a通道（红-绿）
            lab_float[:, :, 1] -= np.random.uniform(2, 8)   # b通道（蓝-黄）
        else:
            # 向冷色偏移：增加蓝色，减少红色
            lab_float[:, :, 2] -= np.random.uniform(5, 15)  # a通道
            lab_float[:, :, 1] += np.random.uniform(2, 8)   # b通道
        
        # 限制在有效范围内
        lab_float[:, :, 1:] = np.clip(lab_float[:, :, 1:], 0, 255)
        lab_shifted = lab_float.astype(np.uint8)
        
        return cv2.cvtColor(lab_shifted, cv2.COLOR_LAB2BGR)
    
    def apply_comprehensive_degradation(self, hr_image: np.ndarray) -> np.ndarray:
        """应用综合退化 - 模拟真实LED显示场景"""
        degraded = hr_image.copy()
        
        # 基于实际工程经验选择退化组合
        # 低灰偏色 + 均匀性问题是最常见的组合
        base_degradations = ['low_gray_color_shift', 'non_uniform_brightness']
        
        # 随机添加其他退化
        additional_degradations = ['motion_tailing', 'ic_coupling_artifacts', 
                                 'bit_depth_quantization', 'color_temperature_shift']
        num_additional = np.random.randint(1, 3)
        selected_additional = np.random.choice(additional_degradations, num_additional, replace=False)
        
        all_degradations = base_degradations + selected_additional.tolist()
        
        # 按顺序应用退化（模拟实际信号处理链路）
        for deg_type in all_degradations:
            if deg_type == 'low_gray_color_shift':
                current_level = np.random.choice(self.config['low_gray']['current_levels'])
                degraded = self.low_gray_color_shift(degraded, current_level)
            elif deg_type in self.degradation_functions:
                degraded = self.degradation_functions[deg_type](degraded)
        
        # 最后进行分辨率降级（模拟物理显示限制）
        h, w = degraded.shape[:2]
        scale_factor = np.random.choice([2, 4])  # 2倍或4倍超分
        lr = cv2.resize(degraded, (w//scale_factor, h//scale_factor), 
                       interpolation=cv2.INTER_CUBIC)
        final_degraded = cv2.resize(lr, (w, h), interpolation=cv2.INTER_CUBIC)
        
        return final_degraded, scale_factor

def create_professional_dataset():
    """创建专业级LED显示数据集"""
    simulator = ProfessionalLEDDegradation()
    
    # 创建配置文件夹
    import os
    os.makedirs("configs", exist_ok=True)
    
    # 保存工程配置
    with open("configs/engineering_config.json", 'w') as f:
        json.dump(simulator.config, f, indent=2)
    
    # 使用高质量图像源
    input_dir = "datasets/high_quality_source"
    output_dir = "datasets/professional_led_dataset"
    
    os.makedirs(os.path.join(output_dir, "HR"), exist_ok=True)
    os.makedirs(os.path.join(output_dir, "LR"), exist_ok=True)
    os.makedirs(os.path.join(output_dir, "metadata"), exist_ok=True)
    
    # 收集图像文件
    image_files = []
    for root, dirs, files in os.walk(input_dir):
        for file in files:
            if file.lower().endswith(('.png', '.jpg', '.jpeg')):
                image_files.append(os.path.join(root, file))
    
    print(f"找到 {len(image_files)} 个高质量源图像")
    
    # 为每个图像生成多个退化版本
    metadata = []
    for img_path in image_files[:200]:  # 处理200个图像，每个生成2-3个退化版本
        hr_image = cv2.imread(img_path)
        if hr_image is None:
            continue
        
        base_name = os.path.splitext(os.path.basename(img_path))[0]
        
        # 为每个源图像生成多个退化版本
        num_variations = np.random.randint(2, 4)
        for variation in range(num_variations):
            # 应用综合退化
            lr_image, scale_factor = simulator.apply_comprehensive_degradation(hr_image)
            
            # 保存图像对
            output_name = f"{base_name}_v{variation+1}"
            cv2.imwrite(os.path.join(output_dir, "HR", f"{output_name}.png"), hr_image)
            cv2.imwrite(os.path.join(output_dir, "LR", f"{output_name}.png"), lr_image)
            
            # 记录元数据 - 确保所有值都是JSON可序列化的Python原生类型
            metadata.append({
                'hr_path': f"HR/{output_name}.png",
                'lr_path': f"LR/{output_name}.png", 
                'scale_factor': int(scale_factor),  # 确保是Python int类型
                'source_image': str(base_name),     # 确保是字符串类型
                'variation': int(variation + 1)     # 确保是Python int类型
            })
    
    # 保存元数据
    with open(os.path.join(output_dir, "metadata", "dataset_info.json"), 'w') as f:
        json.dump(metadata, f, indent=2)
    
    print(f"专业数据集创建完成，共生成 {len(metadata)} 个图像对")
    return output_dir

if __name__ == "__main__":
    dataset_path = create_professional_dataset()
    print(f"数据集保存在: {dataset_path}")